package cz.cuni.amis.episodic.bayes.memories;

import cz.cuni.amis.episodic.bayes.utils.NodeSnapshot;
import cz.cuni.amis.episodic.bayes.utils.Utils;
import cz.cuni.amis.episodic.bayes.utils.chart.GraphPaintingDevice;
import cz.cuni.amis.episodic.dybanem.DBN;
import cz.cuni.amis.episodic.dybanem.Evidence;

import java.util.*;

import smile.Network;

/**
 * Finds mem that minimizes KL from (prior + mem) to real distribution.
 *
 * @author ik
 */
public class MinimizeKLMemoryCreator extends KLBasedMemoryCreator {

    public MinimizeKLMemoryCreator(GraphPaintingDevice graphPaintingDevice) {
        super(graphPaintingDevice, "retro_ig");
    }

    @Override
    public MemoryElement createMemoryInternal(String[] targets, NodeSnapshot[] trueSnapshots, DBN dbn, Collection<? extends Evidence> evidences, int iteration) {
        // find possible interval mems
        List<MemoryElement> mems = new ArrayList<>();
        for (int i = 0; i < trueSnapshots.length; i++) {
            mems.addAll(computeSameOutcomeSequences(trueSnapshots[i].getProbability(), trueSnapshots[i].getNode(), dbn.getNet()));
        }

        // create snapshot given current memory
        // NodeSnapshot[] currentSnapshots = dbn.getNodesSnapshot(targets);
        // compute KL from current knowledge to the real state
        // double[][] kl = Utils.computeKL(currentSnapshots, trueSnapshots);
        double[][] kl = Utils.computeKL(trueSnapshots, currentSnapshots);

        // compute information gain obtained by remembering the mem
        for (MemoryElement mem : mems) {
            // set evidence from previous iterations
            dbn.resetEvidence(evidences);
            // set evidence represented by mem
            mem.setEvidence(dbn);
            
            // compute KL from updated knowledge to the real state
            dbn.getNet().updateBeliefs();
            NodeSnapshot[] updatedSnapshots = dbn.getNodesSnapshot(targets);
            //double[][] klUpdated = Utils.computeKL(updatedSnapshots, trueSnapshots);
            double[][] klUpdated = Utils.computeKL(trueSnapshots, updatedSnapshots);

            // compute difference of kl and klUpdated and set it as an information gain of the mem
            double sum = 0;
            for (int i = 0; i < kl.length; i++) {
                for (int j = 0; j < kl[0].length; j++) {
                    sum += kl[i][j] - klUpdated[i][j];
                }
            }
            mem.setSurprise(sum);

            // clear the evidence
            dbn.resetEvidence(evidences);
        }

        // find the mem with maximal information gain
        MemoryElement mem = Collections.max(mems, new Comparator<MemoryElement>() {

            @Override
            public int compare(MemoryElement o1, MemoryElement o2) {
                return (int) Math.signum(o1.surprise - o2.surprise);
            }
        });

        // TODO debug
        double[] gains = new double[mems.size()];
        int ii = 0;
        for(MemoryElement m : mems) gains[ii++] = m.surprise;
        
        // set evidence of best mem
        mem.setEvidence(dbn);

        // draw graph
        if (graphPaintingDevice != null) {

            // find index of the best node
            int nodeIx = 0;
            for (int i = 0; i < targets.length; i++) {
                if (targets[i].equals(mem.getNode())) {
                    nodeIx = i;
                    break;
                }
            }
            mem.markInKLGraph(kl[nodeIx]); // we set this only because of drawing, NaN has specal color and here it indicates the max value        
            drawKLGraph(kl, targets, iteration);
        }
        
        return mem;
    }

    /**
     * Transform sequence of prob. distributions into intervals where the most
     * probable outcome is the same.
     *
     * @param probability
     * @return
     */
    List<MemoryElement> computeSameOutcomeSequences(double[][] probability, String node, Network net) {
        List<MemoryElement> l = new ArrayList<>();
        int start = 0;
        int lastMaxIx = Utils.maxIndex(probability[0]);
        final int N = probability[0].length;
        for (int i = 0; i < probability.length; i++) {
            int maxIx = Utils.maxIndex(probability[i]);
            if (maxIx != lastMaxIx) {
                // store actual sequence
                MemoryElement interval = new MemoryElement(node,
                        start, i - 1, 0d, net.getOutcomeIds(node)[lastMaxIx]);
                l.add(interval);
                // start new sequence
                start = i;
                lastMaxIx = maxIx;
            }
        }
        l.add(new MemoryElement(node,
                start, probability.length - 1, 0d, net.getOutcomeIds(node)[lastMaxIx]));
        return l;
    }
}
